import os
import base64
import hashlib
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from django.conf import settings
from django.core.cache import cache
from django.utils import timezone
import logging

logger = logging.getLogger(__name__)


class EncryptionManager:
    """加密管理器"""
    
    def __init__(self):
        self._fernet = None
        self._init_encryption()
    
    def _init_encryption(self):
        """初始化加密器"""
        try:
            # 从设置中获取密钥，如果没有则生成新的
            secret_key = getattr(settings, 'ENCRYPTION_SECRET_KEY', None)
            if not secret_key or secret_key == 'your-secret-key-here':
                # 生成新的密钥
                secret_key = Fernet.generate_key()
                logger.warning("未找到加密密钥，已生成新密钥。请将其保存到设置中。")
                logger.warning(f"ENCRYPTION_SECRET_KEY = '{secret_key.decode()}'")

            if isinstance(secret_key, str):
                # 如果是字符串，尝试解码为base64
                try:
                    secret_key = secret_key.encode()
                    # 验证是否为有效的Fernet密钥
                    Fernet(secret_key)
                except Exception:
                    # 如果不是有效的Fernet密钥，生成新的
                    secret_key = Fernet.generate_key()
                    logger.warning("无效的加密密钥，已生成新密钥。")
                    logger.warning(f"ENCRYPTION_SECRET_KEY = '{secret_key.decode()}'")

            self._fernet = Fernet(secret_key)
        except Exception as e:
            logger.error(f"初始化加密器失败: {e}")
            # 使用默认密钥作为后备
            default_key = Fernet.generate_key()
            self._fernet = Fernet(default_key)
            logger.warning(f"使用临时密钥: {default_key.decode()}")
    
    def encrypt(self, data: str) -> str:
        """加密数据"""
        try:
            if not data:
                return data
            
            encrypted_data = self._fernet.encrypt(data.encode())
            return base64.b64encode(encrypted_data).decode()
        except Exception as e:
            logger.error(f"数据加密失败: {e}")
            raise
    
    def decrypt(self, encrypted_data: str) -> str:
        """解密数据"""
        try:
            if not encrypted_data:
                return encrypted_data
            
            encrypted_bytes = base64.b64decode(encrypted_data.encode())
            decrypted_data = self._fernet.decrypt(encrypted_bytes)
            return decrypted_data.decode()
        except Exception as e:
            logger.error(f"数据解密失败: {e}")
            raise
    
    def hash_password(self, password: str, salt: str = None) -> tuple:
        """哈希密码"""
        if salt is None:
            salt = os.urandom(32)
        elif isinstance(salt, str):
            salt = salt.encode()
        
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = kdf.derive(password.encode())
        return base64.b64encode(key).decode(), base64.b64encode(salt).decode()
    
    def verify_password(self, password: str, hashed_password: str, salt: str) -> bool:
        """验证密码"""
        try:
            new_hash, _ = self.hash_password(password, salt)
            return new_hash == hashed_password
        except Exception:
            return False


class SecurityValidator:
    """安全验证器"""
    
    @staticmethod
    def validate_ssh_key(private_key: str) -> bool:
        """验证SSH私钥格式"""
        try:
            import paramiko
            from io import StringIO
            
            # 尝试加载私钥
            paramiko.RSAKey.from_private_key(StringIO(private_key))
            return True
        except Exception:
            try:
                # 尝试其他密钥类型
                paramiko.DSSKey.from_private_key(StringIO(private_key))
                return True
            except Exception:
                try:
                    paramiko.ECDSAKey.from_private_key(StringIO(private_key))
                    return True
                except Exception:
                    return False
    
    @staticmethod
    def validate_ip_address(ip: str) -> bool:
        """验证IP地址格式"""
        import ipaddress
        try:
            ipaddress.ip_address(ip)
            return True
        except ValueError:
            return False
    
    @staticmethod
    def validate_port(port: int) -> bool:
        """验证端口号"""
        return 1 <= port <= 65535
    
    @staticmethod
    def is_safe_command(command: str) -> bool:
        """检查命令是否安全"""
        dangerous_commands = [
            'rm -rf', 'mkfs', 'dd if=', 'format', 'fdisk',
            'shutdown', 'reboot', 'halt', 'poweroff',
            'passwd', 'userdel', 'usermod',
            'iptables -F', 'ufw --force',
            'chmod 777', 'chown -R'
        ]
        
        command_lower = command.lower().strip()
        for dangerous in dangerous_commands:
            if dangerous in command_lower:
                return False
        
        return True


class RateLimiter:
    """速率限制器"""
    
    def __init__(self, max_requests: int = 100, window_seconds: int = 3600):
        self.max_requests = max_requests
        self.window_seconds = window_seconds
    
    def is_allowed(self, identifier: str) -> bool:
        """检查是否允许请求"""
        cache_key = f"rate_limit:{identifier}"
        current_requests = cache.get(cache_key, 0)
        
        if current_requests >= self.max_requests:
            return False
        
        # 增加计数
        cache.set(cache_key, current_requests + 1, self.window_seconds)
        return True
    
    def get_remaining(self, identifier: str) -> int:
        """获取剩余请求次数"""
        cache_key = f"rate_limit:{identifier}"
        current_requests = cache.get(cache_key, 0)
        return max(0, self.max_requests - current_requests)


class AuditLogger:
    """审计日志记录器"""
    
    @staticmethod
    def log_ssh_connection(user, host, action, result, details=None):
        """记录SSH连接日志"""
        log_data = {
            'user': user.username if user else 'anonymous',
            'user_id': user.id if user else None,
            'host': host.hostname if host else 'unknown',
            'host_ip': host.ip if host else 'unknown',
            'action': action,
            'result': result,
            'timestamp': timezone.now().isoformat(),
            'details': details or {}
        }
        
        logger.info(f"SSH_AUDIT: {log_data}")
        
        # 可以扩展到数据库记录
        try:
            from .models import AuditLog
            AuditLog.objects.create(
                user=user,
                action=action,
                resource_type='ssh_connection',
                resource_id=host.id if host else None,
                result=result,
                details=log_data
            )
        except Exception as e:
            logger.error(f"保存审计日志失败: {e}")
    
    @staticmethod
    def log_command_execution(session, command, result, exit_code=None):
        """记录命令执行日志"""
        log_data = {
            'session_id': str(session.session_id),
            'user': session.user.username,
            'host': session.host.hostname,
            'command': command,
            'result': result,
            'exit_code': exit_code,
            'timestamp': timezone.now().isoformat(),
            'is_dangerous': not SecurityValidator.is_safe_command(command)
        }
        
        logger.info(f"COMMAND_AUDIT: {log_data}")
        
        # 危险命令特别记录
        if log_data['is_dangerous']:
            logger.warning(f"DANGEROUS_COMMAND: {log_data}")


class ConnectionPoolManager:
    """连接池管理器"""
    
    def __init__(self, max_connections: int = 50, max_idle_time: int = 300):
        self.max_connections = max_connections
        self.max_idle_time = max_idle_time
        self._connections = {}
        self._connection_stats = {}
    
    def get_connection_key(self, host, user):
        """生成连接键"""
        return f"{host.ip}:{host.ssh_port}:{user}"
    
    def get_connection(self, host, user):
        """获取连接"""
        key = self.get_connection_key(host, user)
        
        # 检查是否有可用连接
        if key in self._connections:
            connection = self._connections[key]
            if self._is_connection_valid(connection):
                self._update_connection_stats(key, 'reused')
                return connection
            else:
                # 清理无效连接
                self._remove_connection(key)
        
        # 检查连接池是否已满
        if len(self._connections) >= self.max_connections:
            self._cleanup_idle_connections()
        
        # 创建新连接
        connection = self._create_connection(host, user)
        if connection:
            self._connections[key] = connection
            self._update_connection_stats(key, 'created')
        
        return connection
    
    def _is_connection_valid(self, connection):
        """检查连接是否有效"""
        try:
            # 检查连接是否还活着
            transport = connection.get_transport()
            return transport and transport.is_active()
        except Exception:
            return False
    
    def _create_connection(self, host, user):
        """创建新连接"""
        try:
            import paramiko
            
            ssh_client = paramiko.SSHClient()
            ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # 连接配置
            connect_kwargs = {
                'hostname': host.ip,
                'port': host.ssh_port,
                'username': host.ssh_username,
                'timeout': 10
            }
            
            # 认证方式
            if host.ssh_auth_type == 'password':
                password = host.decrypt_password()
                connect_kwargs['password'] = password
            else:
                private_key = host.decrypt_private_key()
                if private_key:
                    from io import StringIO
                    key = paramiko.RSAKey.from_private_key(StringIO(private_key))
                    connect_kwargs['pkey'] = key
                else:
                    raise Exception("SSH私钥未配置")
            
            ssh_client.connect(**connect_kwargs)
            return ssh_client
            
        except Exception as e:
            logger.error(f"创建SSH连接失败: {e}")
            return None
    
    def _remove_connection(self, key):
        """移除连接"""
        if key in self._connections:
            try:
                self._connections[key].close()
            except Exception:
                pass
            del self._connections[key]
        
        if key in self._connection_stats:
            del self._connection_stats[key]
    
    def _cleanup_idle_connections(self):
        """清理空闲连接"""
        from django.utils import timezone
        current_time = timezone.now()
        
        keys_to_remove = []
        for key, stats in self._connection_stats.items():
            if (current_time - stats['last_used']).seconds > self.max_idle_time:
                keys_to_remove.append(key)
        
        for key in keys_to_remove:
            self._remove_connection(key)
    
    def _update_connection_stats(self, key, action):
        """更新连接统计"""
        from django.utils import timezone
        
        if key not in self._connection_stats:
            self._connection_stats[key] = {
                'created_at': timezone.now(),
                'last_used': timezone.now(),
                'use_count': 0
            }
        
        stats = self._connection_stats[key]
        stats['last_used'] = timezone.now()
        stats['use_count'] += 1
    
    def get_stats(self):
        """获取连接池统计"""
        return {
            'total_connections': len(self._connections),
            'max_connections': self.max_connections,
            'connection_details': [
                {
                    'key': key,
                    'created_at': stats['created_at'],
                    'last_used': stats['last_used'],
                    'use_count': stats['use_count']
                }
                for key, stats in self._connection_stats.items()
            ]
        }
    
    def close_all(self):
        """关闭所有连接"""
        for key in list(self._connections.keys()):
            self._remove_connection(key)


# 全局实例
encryption_manager = EncryptionManager()
security_validator = SecurityValidator()
ssh_rate_limiter = RateLimiter(max_requests=50, window_seconds=300)  # 5分钟50次
connection_pool = ConnectionPoolManager()
audit_logger = AuditLogger()


# 装饰器
def rate_limit(identifier_func=None, max_requests=100, window_seconds=3600):
    """速率限制装饰器"""
    def decorator(func):
        def wrapper(request, *args, **kwargs):
            if identifier_func:
                identifier = identifier_func(request)
            else:
                identifier = request.user.username if request.user.is_authenticated else request.META.get('REMOTE_ADDR')
            
            limiter = RateLimiter(max_requests, window_seconds)
            if not limiter.is_allowed(identifier):
                from django.http import JsonResponse
                return JsonResponse({
                    'success': False,
                    'message': '请求过于频繁，请稍后再试'
                }, status=429)
            
            return func(request, *args, **kwargs)
        return wrapper
    return decorator


def audit_log(action, resource_type=None):
    """审计日志装饰器"""
    def decorator(func):
        def wrapper(request, *args, **kwargs):
            start_time = timezone.now()
            result = 'success'
            error_message = None
            
            try:
                response = func(request, *args, **kwargs)
                return response
            except Exception as e:
                result = 'error'
                error_message = str(e)
                raise
            finally:
                # 记录审计日志
                try:
                    from .models import AuditLog
                    AuditLog.objects.create(
                        user=request.user if request.user.is_authenticated else None,
                        action=action,
                        resource_type=resource_type or 'unknown',
                        resource_id=kwargs.get('pk'),
                        result=result,
                        details={
                            'method': request.method,
                            'path': request.path,
                            'duration': (timezone.now() - start_time).total_seconds(),
                            'error': error_message
                        }
                    )
                except Exception as e:
                    logger.error(f"记录审计日志失败: {e}")
        
        return wrapper
    return decorator
